package cn.qiao.mdc.thread;

import cn.qiao.mdc.consts.Constants;
import org.slf4j.MDC;

import java.util.Map;
import java.util.concurrent.Callable;

public class CallableWrapper<V> implements Callable<V> {

    private final Callable<V> callable;
    private final Map<String, String> contextMap;

    public CallableWrapper(Callable<V> callable) {
        this.callable = callable;
        this.contextMap = MDC.getCopyOfContextMap();
    }

    @Override
    public V call() throws Exception {
        try {
            if (contextMap != null) {
                MDC.setContextMap(contextMap);
            } else {
                MDC.put(Constants.TRACE_ID, Constants.getTraceId());
            }
            return callable != null ? callable.call() : null;
        } finally {
            MDC.clear();
        }
    }
}
